from io import BytesIO
from typing import List

from fastapi import APIRouter, Body
from openpyxl.workbook import Workbook
from starlette.responses import StreamingResponse
from tortoise.exceptions import DoesNotExist

from app.schemas.electric_price_info import ElectricPriceInfoCreate, ElectricPriceInfoOut, PriceChartOut, \
    ElectricPriceInfoDetailOut, ElectricPriceInfoCreateList
from app.crud import electric_price_info as crud_price_info
from app.utils.Result import Result

router = APIRouter()

@router.post(
    "/electric-price-info/",
    response_model=Result[bool],
    summary="创建电价信息及对应电价数值",
    description="根据传入的电价信息及一组电价数值，创建电价主表和电价值子表数据。"
)
async def create_price_info_api(data: ElectricPriceInfoCreate):
    success = await crud_price_info.create_price_info(data)
    if not success:
        return Result.error(code=400, msg="电价数值条目数量与划分粒度不符")
    return Result.ok(msg=f"成功添加 {len(data.values)} 个电价数值")

@router.post(
    "/electric-price-info/create-list",
    response_model=Result[bool],
    summary="创建电价信息及对应电价数值",
    description="根据传入的电价信息及一组电价数值，创建电价主表和电价值子表数据。"
)
async def create_price_info_api(data: ElectricPriceInfoCreateList):
    success = await crud_price_info.create_price_info_list(data)
    if not success:
        return Result.error(code=400, msg="电价数值条目数量与划分粒度不符")
    return Result.ok(msg=f"成功添加 {len(data.values)} 个电价数值")

@router.delete(
    "/electric-price-info/batch",
    response_model=Result[str],
    summary="批量删除电价信息及其电价数值",
    description="根据电价信息ID列表，删除多个电价主表记录及其关联的电价值记录。"
)
async def delete_price_infos_api(
        price_ids: List[int] = Body(..., embed=True, description="要删除的电价信息ID列表")
):
    deleted_count = await crud_price_info.delete_price_infos(price_ids)
    if deleted_count == 0:
        return Result.error(code=404, msg="未找到匹配的电价信息")
    return Result.ok(msg=f"删除成功，共删除 {deleted_count} 条记录")

@router.delete(
    "/electric-price-info/{price_id}",
    response_model=Result[str],
    summary="删除电价信息及其对应的电价数值",
    description="根据电价信息ID，删除电价主表记录及其关联的电价值记录。"
)
async def delete_price_info_api(price_id: int):
    success = await crud_price_info.delete_price_info(price_id)
    if not success:
        return Result.error(code=404, msg="未找到对应的电价信息")
    return Result.ok(msg="删除成功")

@router.get(
    "/electric-price-info/",
    response_model=Result[List[ElectricPriceInfoOut]],
    summary="获取所有电价主表信息",
    description="获取所有 ElectricPriceInfo 记录（不包含子表 ElectricPriceValue）。"
)
async def get_all_price_infos():
    result = await crud_price_info.list_price_info()
    return Result.ok(data=result)

@router.post(
    "/parameters/detail/by_ids",
    response_model=Result[List[ElectricPriceInfoDetailOut]],
    summary="根据ID列表获取电价详细信息",
    description="传入参数ID列表，查询对应的电价详细信息"
)
async def list_pricess_detail_by_ids_api(price_ids: List[int] = Body(..., embed=True)):
    try:
        price = await crud_price_info.list_prices_detail_by_ids(price_ids)
        return Result.ok(data=price)
    except Exception as e:
        print(e)
        return Result.error(code=500, msg="内部服务器错误")


@router.get(
    "/electric-price-info/{price_id}/values",
    response_model=Result[PriceChartOut],
    summary="根据 price_id 获取电价数值列表",
    description="返回 time_list 和 value_list，用于图表展示"
)
async def get_price_values(price_id: int):
    result = await crud_price_info.list_price_values_by_price_id(price_id)
    return Result.ok(data=result)


@router.get(
    "/electric-price-info/category/{price_type}",
    response_model=Result[List[ElectricPriceInfoOut]],
    summary="根据类别ID获取电价列表",
    description="根据参数类别ID查询该类别下所有电价信息"
)
async def list_parameters_by_category_api(price_type: int):
    try:
        params = await crud_price_info.list_parameters_by_category(price_type)
        return Result.ok(data=params)
    except DoesNotExist:
        return Result.error(code=404, msg="未找到对应的电价类别")


@router.post(
    "/parameters/detail/export_prices_by_ids",
    summary="根据ID列表导出电价详细信息",
    description="传入电价ID列表，导出对应的电价详细信息Excel文件"
)
async def export_prices_detail_by_ids_api(price_ids: List[int] = Body(..., embed=True)):
    try:
        # 从crud层拿数据，返回示例结构类似你给的那样
        prices = await crud_price_info.list_prices_detail_by_ids(price_ids)

        wb = Workbook()
        ws = wb.active
        ws.title = "电价详情"

        # 表头
        ws.append(["电价类别ID", "价格名称", "时间", "价格值"])

        for price in prices:
            price_type = price.get("price_type")
            price_name = price.get("price_name")
            value_list = price.get("value_list", [])
            for time_str, val in zip(price.get("time_list"), price.get("value_list")):
                ws.append([
                    price_type,
                    price_name,
                    time_str.replace(tzinfo=None).isoformat(),
                    val
                ])

        stream = BytesIO()
        wb.save(stream)
        stream.seek(0)

        return StreamingResponse(
            stream,
            media_type="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
            headers={"Content-Disposition": "attachment; filename=prices_export.xlsx"}
        )
    except Exception as e:
        print(f"导出电价详细信息时发生错误: {e}")
        return StreamingResponse(
            BytesIO("导出失败".encode("utf-8")),
            media_type="text/plain",
            status_code=500
        )